import torch
import numpy as np
from PIL import Image
from torchvision import transforms
from dml import TripletModel  # 假设 TripletModel 定义在 dml_torch.dml 模块中

# 加载训练好的模型
model = TripletModel()
model.load_state_dict(torch.load('src/model_pth/triplet_model_iu_xray.pth',weights_only=False))
model.eval()

# 定义图像预处理
preprocess = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.Grayscale(num_output_channels=1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

def get_embedding(image_path):
    # 加载图像并进行预处理
    image = Image.open(image_path)
    image = preprocess(image).unsqueeze(0)  # 增加 batch 维度
    
    # 获取嵌入向量
    with torch.no_grad():
        embedding = model(image)
    
    # 打印嵌入向量
    print("Embedding Vector:")
    print(embedding.numpy())
    
    return embedding.numpy()

# embeddings using


# 示例调用
if __name__ == "__main__":
    image_path = 'README/image.png'  # 替换为你的图片路径
    embedding = get_embedding(image_path)
    print(embedding)